#!/bin/bash
set -euo pipefail

# 脚本名称：resize_vdisk.sh

# 功能说明：
# 本脚本用于扩展虚拟磁盘，并同步调整分区和文件系统大小。

# 主要流程包括：
# 1. 安装必要的依赖工具（如 qemu-img、qemu-nbd、parted 等）。
# 2. 选择可用的 NBD 设备，并将虚拟磁盘挂载到 NBD 设备。
# 3. 预检磁盘文件，显示当前分区信息，并提供用户确认。
# 4. 通过 qemu-img 扩展磁盘容量，并重新附加 NBD 设备。
# 5. 调整最后一个分区的大小，并扩展 ext2/3/4 文件系统以匹配新容量。
# 6. 断开 NBD 设备，完成操作。

# 注意事项：
# - 仅支持 ext2/3/4 文件系统的分区扩展。
# - 需要 root 权限运行部分命令，如 qemu-nbd、parted、resize2fs。
# - 在执行磁盘扩展前，建议备份重要数据，以防万一。

#--------- 依赖管理 ---------#
install_dependencies() {
    local required=(qemu-img qemu-nbd partprobe resize2fs parted lsblk blkid)
    local missing=()

    for cmd in "${required[@]}"; do
        if ! command -v "$cmd" &>/dev/null; then
            missing+=("$cmd")
        fi
    done

    if [ ${#missing[@]} -gt 0 ]; then
        echo "缺少依赖: ${missing[*]}，正在安装..."
        sudo apt-get update -qq
        # 安装提供相应命令的软件包
        sudo apt-get install -y -qq qemu-utils parted e2fsprogs util-linux >/dev/null
    fi
}

#--------- 全局变量 ---------#
nbd_device=""
virtual_disk_attached=false

# 清理函数：断开当前连接的 NBD 设备
cleanup() {
    if $virtual_disk_attached && [ -n "$nbd_device" ]; then
        echo "断开设备: $nbd_device"
        sudo qemu-nbd -d "$nbd_device" >/dev/null
        virtual_disk_attached=false
    fi
}
trap cleanup EXIT SIGINT SIGTERM

#--------- NBD设备管理 ---------#
# 从 /dev/nbd* 中选择一个空闲的设备
select_nbd() {
    for dev in /dev/nbd*; do
        # 检查设备是否存在对应的 pid 文件且内容为空（表示未被占用）
        pid_file="/sys/block/$(basename "$dev")/pid"
        if [ -f "$pid_file" ]; then
            if [ ! -s "$pid_file" ]; then
                echo "$dev"
                return 0
            fi
        else
            # 若 pid 文件不存在，则认为设备可用
            echo "$dev"
            return 0
        fi
    done
    return 1
}

# 附加虚拟磁盘到空闲的 nbd 设备
attach_disk() {
    local disk_image="$1"
    sudo modprobe nbd max_part=8
    nbd_device=$(select_nbd)
    if [ -z "$nbd_device" ]; then
        echo "无可用NBD设备" >&2
        exit 1
    fi

    echo "正在附加 $disk_image 到 $nbd_device..."
    sudo qemu-nbd -c "$nbd_device" "$disk_image" || { echo "附加失败"; exit 1; }
    sudo udevadm settle
    virtual_disk_attached=true
}

#--------- 磁盘操作 ---------#
# 验证磁盘文件是否存在，不再显示 qemu-img info 信息
validate_disk() {
    local disk_image="$1"
    if [ ! -f "$disk_image" ]; then
        echo "无效磁盘文件: $disk_image" >&2
        exit 1
    fi
    echo "磁盘验证通过"
}

# 设置虚拟磁盘的目标大小
resize_disk() {
    local disk_image="$1"
    local size
    local size_re='^[0-9]+[GM]$'
    read -rp "请输入虚拟磁盘扩展后的目标大小（例如：'20G' 或 '500M'）: " size
    if [[ ! "$size" =~ $size_re ]]; then
        echo "格式错误，请输入类似 '20G' 或 '500M'" >&2
        exit 1
    fi

    echo "将虚拟磁盘 $disk_image 扩展到目标大小 $size ..."
    sudo qemu-img resize "$disk_image" "$size" || { echo "扩展失败"; exit 1; }
}

#--------- 分区及文件系统操作 ---------#
# 验证最后一个分区的文件系统是否为 ext2/3/4
validate_filesystem() {
    local nbd="$1"
    local part
    part="/dev/$(lsblk -nlo NAME,TYPE "$nbd" | awk '$2=="part" {print $1}' | tail -n1)"
    if [ -z "$part" ]; then
        echo "未找到分区" >&2
        exit 1
    fi

    local fs_type
    fs_type=$(sudo blkid -s TYPE -o value "$part")
    case "$fs_type" in
        ext2|ext3|ext4)
            ;;
        *)
            echo "不支持的文件系统: $fs_type" >&2
            exit 1
            ;;
    esac
}

# 调整最后一个分区并扩展文件系统
resize_partition() {
    local nbd="$1"
    local part_info part_num part_path

    part_info=$(lsblk -nlo NAME,TYPE "$nbd" | awk '$2=="part" {print $1}' | tail -n1)
    if [ -z "$part_info" ]; then
        echo "未找到分区" >&2
        exit 1
    fi

    part_num=$(echo "$part_info" | grep -o '[0-9]*$')
    echo "调整分区 #$part_num"
    sudo parted -s "$nbd" resizepart "$part_num" 100% || exit 1
    sudo partprobe "$nbd"
    part_path="/dev/$(lsblk -nlo NAME,TYPE "$nbd" | awk '$2=="part" {print $1}' | tail -n1)"
    sudo resize2fs "$part_path" || { echo "文件系统调整失败"; exit 1; }
}

#--------- 主流程 ---------#
main() {
    install_dependencies

    local disk_image="${1:-}"
    if [ -z "$disk_image" ]; then
        read -rp "输入磁盘路径: " disk_image
    fi

    # 预检阶段：验证并附加磁盘显示分区布局
    validate_disk "$disk_image"
    attach_disk "$disk_image"
    echo -e "\n当前分区布局:"
    lsblk "$nbd_device"
    read -rp "确认要继续? [y/N] " confirm
    if [[ ! "$confirm" =~ ^[yY]$ ]]; then
        echo "操作取消"
        exit 0
    fi
    cleanup  # 断开预检阶段的连接

    # 扩展阶段：调整磁盘、重新附加、验证及扩展分区和文件系统
    resize_disk "$disk_image"
    attach_disk "$disk_image"
    validate_filesystem "$nbd_device"
    resize_partition "$nbd_device"
    
    echo -e "\n最终分区状态:"
    lsblk "$nbd_device"
    echo "操作成功完成"
}

main "$@"
